import os
import sys
import time
import spidev
import numpy as np
from PIL import Image, ImageDraw, ImageFont, ImageOps

LCD_WIDTH  = 320
LCD_HEIGHT = 240
SPI_CHUNK_ROWS = 6

import lgpio

LCD_RST_PIN = 22
LCD_DC_PIN  = 25
LCD_BL_PIN  = 24

h = lgpio.gpiochip_open(0)

class LCD_2Inch:
    def __init__(self):
        lgpio.gpio_claim_output(h, LCD_DC_PIN, 0)
        lgpio.gpio_claim_output(h, LCD_RST_PIN, 1)
        lgpio.gpio_claim_output(h, LCD_BL_PIN, 0)
        lgpio.tx_pwm(h, LCD_BL_PIN, 1000, 90)
        self.spi = spidev.SpiDev()
        self.spi.open(0, 0)
        self.spi.max_speed_hz = 8000000
        self.spi.mode = 0b00
        self.lcd_init()

    def write_cmd(self, cmd):
        lgpio.gpio_write(h, LCD_DC_PIN, 0)
        self.spi.writebytes([cmd])

    def write_data(self, value):
        lgpio.gpio_write(h, LCD_DC_PIN, 1)
        if isinstance(value, int):
            self.spi.writebytes([value])
        else:
            self.spi.writebytes(value)

    def reset(self):
        lgpio.gpio_write(h, LCD_RST_PIN, 1); time.sleep(0.001)
        lgpio.gpio_write(h, LCD_RST_PIN, 0); time.sleep(0.080)
        lgpio.gpio_write(h, LCD_RST_PIN, 1); time.sleep(0.120)

    def lcd_init(self):
        self.reset()
        self.write_cmd(0x01); time.sleep(0.12)
        self.write_cmd(0x11); time.sleep(0.12)
        self.write_cmd(0xCF); self.write_data([0x00, 0xC1, 0x30])
        self.write_cmd(0xED); self.write_data([0x64, 0x03, 0x12, 0x81])
        self.write_cmd(0xE8); self.write_data([0x85, 0x00, 0x79])
        self.write_cmd(0xCB); self.write_data([0x39, 0x2C, 0x00, 0x34, 0x02])
        self.write_cmd(0xF7); self.write_data(0x20)
        self.write_cmd(0xEA); self.write_data([0x00, 0x00])
        self.write_cmd(0xC0); self.write_data(0x1D)
        self.write_cmd(0xC1); self.write_data(0x12)
        self.write_cmd(0xC5); self.write_data([0x33, 0x3F])
        self.write_cmd(0xC7); self.write_data(0x92)
        self.write_cmd(0x3A); self.write_data(0x55)
        self.write_cmd(0x36); self.write_data(0x28)  # landscape + BGR
        self.write_cmd(0xB1); self.write_data([0x00, 0x12])
        self.write_cmd(0xB6); self.write_data([0x0A, 0xA2])
        self.write_cmd(0x44); self.write_data(0x02)
        self.write_cmd(0xF2); self.write_data(0x00)
        self.write_cmd(0x26); self.write_data(0x01)
        self.write_cmd(0xE0); self.write_data([0x0F,0x22,0x1C,0x1B,0x08,0x0F,0x48,0xB8,0x34,0x05,0x0C,0x09,0x0F,0x07,0x00])
        self.write_cmd(0xE1); self.write_data([0x00,0x23,0x24,0x07,0x10,0x07,0x38,0x47,0x4B,0x0A,0x13,0x06,0x30,0x38,0x0F])
        self.write_cmd(0x11); time.sleep(0.12)
        self.write_cmd(0x29); time.sleep(0.02)

    def display(self, img):
        if img.mode != "RGB":
            img = img.convert("RGB")
        if img.size != (LCD_WIDTH, LCD_HEIGHT):
            img = ImageOps.fit(img, (LCD_WIDTH, LCD_HEIGHT), method=Image.Resampling.LANCZOS)

        self.write_cmd(0x2A)
        self.write_data([0, 0, (LCD_WIDTH-1)>>8, (LCD_WIDTH-1)&0xFF])
        self.write_cmd(0x2B)
        self.write_data([0, 0, (LCD_HEIGHT-1)>>8, (LCD_HEIGHT-1)&0xFF])
        self.write_cmd(0x2C)
        lgpio.gpio_write(h, LCD_DC_PIN, 1)


        arr = np.asarray(img, dtype=np.uint16)
        r = arr[..., 0]
        g = arr[..., 1]
        b = arr[..., 2]
        rgb565 = ((r & 0xF8) << 8) | ((g & 0xFC) << 3) | (b >> 3)
        pixel = np.empty((LCD_HEIGHT, LCD_WIDTH, 2), dtype=np.uint8)
        pixel[..., 0] = (rgb565 >> 8).astype(np.uint8)
        pixel[..., 1] = rgb565.astype(np.uint8)


        for y in range(0, LCD_HEIGHT, SPI_CHUNK_ROWS):
            block = pixel[y:y + SPI_CHUNK_ROWS].tobytes()
            try:
                self.spi.writebytes2(block)
            except AttributeError:
                self.spi.writebytes(list(block))


def make_drawing_page():
    image = Image.new("RGB", (LCD_WIDTH, LCD_HEIGHT), "WHITE")
    draw = ImageDraw.Draw(image)
    try:
        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 22)
        font_small = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 18)
    except Exception:
        font = ImageFont.load_default()
        font_small = font

    draw.text((12, 14), "2.4inch LCD", fill="BLACK", font=font)
    draw.text((12, 50), "SEENGREAT", fill="BLUE", font=font)
    draw.text((184, 50), "RGB565", fill="RED", font=font)

    draw.line((15, 92, 145, 92), fill="BLUE", width=2)
    for x in range(175, 305, 12):
        draw.line((x, 92, min(x+6, 305), 92), fill="RED", width=2)

    draw.ellipse((22, 116, 36, 130), fill="YELLOW")
    draw.rectangle((62, 108, 112, 142), outline="BLUE", width=2)
    draw.rectangle((132, 112, 172, 142), fill="GREEN")
    draw.ellipse((202, 105, 242, 145), outline="MAGENTA", width=2)
    draw.ellipse((266, 109, 298, 141), fill="CYAN")

    draw.text((18, 170), "POINT LINE SHAPE", fill="GREEN", font=font_small)
    draw.text((34, 204), "0123456789", fill="BLACK", font=font)
    return image


def load_picture(path):
    with Image.open(path) as source:
        image = ImageOps.exif_transpose(source).convert("RGB")
        if image.size != (LCD_WIDTH, LCD_HEIGHT):
            image = ImageOps.fit(
                image,
                (LCD_WIDTH, LCD_HEIGHT),
                method=Image.Resampling.LANCZOS,
                centering=(0.5, 0.5),
            )
        return image.copy()


def main():
    lcd = LCD_2Inch()
    picture = [
        "fruits24.bmp",
        "birds.bmp",
        "cat.bmp",
        "fish.bmp",
        "rabbit.bmp",
    ]


    try:
        while True:
            lcd.display(make_drawing_page())
            time.sleep(3)

            for path in picture:
                try:
                    image = load_picture(path)
                except Exception as e:
                    print("Open picture failed:", path, e)
                    image = Image.new("RGB", (LCD_WIDTH, LCD_HEIGHT), "BLACK")
                lcd.display(image)
                time.sleep(3)
    except KeyboardInterrupt:
        print("\nExit")
    finally:
        lgpio.gpiochip_close(h)


if __name__ == "__main__":
    main()
